#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Nov 23 04:06:23 2021

@author: wuzongze
"""

import os

import argparse
import pickle
import numpy as np

import dnnlib
import dnnlib.tflib as tflib
from PIL import Image

def main():
    parser = argparse.ArgumentParser(
        description='combine proj latent codes',
    )
    parser.add_argument('--network',     help='Network pickle filename', required=True)
    parser.add_argument('--source_path',      help='path to source inverted latent codes', required=True)
    parser.add_argument('--target_path',     help='path for saving translated images', required=True)
    
    
    args = parser.parse_args()
    network=args.network
    source_path=args.source_path
    target_path=args.target_path
    
    tflib.init_tf()
    print('Loading networks from "%s"...' % network)
    with dnnlib.util.open_url(network) as fp:
        _, _, Gs = pickle.load(fp)
        
    
    names=os.listdir(source_path)
    names=sorted(names)

    for name in names:
        if name[-4:]=='.npz' :
            tmp=np.load(source_path+name)
            z=tmp['dlatents']
#            print(z.shape)
            img=Gs.run(z,None,output_transform=dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True))
#            print(img.shape)
            name1=name[:-4]+'.jpg'
            
            img=Image.fromarray(img[0]).save(target_path+name1)
            
#%%
if __name__ == "__main__":
    main()




